#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 1;
class KernelEye {
public:
    __aicore__ inline KernelEye() {}
    __aicore__ inline void Init(GM_ADDR y,
                                int16_t num_rows, int16_t num_columns, int32_t num_rows_columns, int8_t formerNum, int16_t tailLength, 
                                TPipe* pipeIn) {
        this->num_rows = num_rows;
        this->num_columns = num_columns;
        this->num_rows_columns = num_rows_columns;

        if (GetBlockIdx() < formerNum)
        {
            this->batch_shape_size = tailLength + 1;
            yGm.SetGlobalBuffer((__gm__ uint32_t*)y + this->num_rows_columns * this->batch_shape_size*2 * GetBlockIdx() + 1);
        }
        else
        {
            this->batch_shape_size = tailLength;
            yGm.SetGlobalBuffer((__gm__ uint32_t*)y +  this->num_rows_columns * this->batch_shape_size*2 * GetBlockIdx() + this->num_rows_columns*2 * formerNum + 1);
        }

        pipeIn->InitBuffer(outQueueY, BUFFER_NUM, this->batch_shape_size * 32);
    }
    __aicore__ inline void Process() {
        LocalTensor<uint32_t> yLocal = outQueueY.AllocTensor<uint32_t>();
        
        Duplicate(yLocal, (uint32_t)0x3FF00000, this->batch_shape_size*8);

        outQueueY.EnQue<uint32_t>(yLocal);
        yLocal = outQueueY.DeQue<uint32_t>();

        DataCopyExtParams copyParams{static_cast<uint16_t>(this->batch_shape_size), 1 * sizeof(uint32_t), 0, static_cast<uint32_t>((this->num_rows_columns)*2 * sizeof(uint32_t) - 4), 0};
        for(int32_t k=0; k<this->num_rows; k++)
        {
            DataCopyPad(yGm[k*(this->num_columns+1)*2], yLocal, copyParams);
        }
        
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe* pipe;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;

    GlobalTensor<uint32_t> yGm;
    int16_t num_rows;
    int16_t num_columns;
    int32_t num_rows_columns;
    int32_t batch_shape_size;
    int16_t interval;
};
class KernelEye_all {
public:
    __aicore__ inline KernelEye_all() {}
    __aicore__ inline void Init(GM_ADDR y,
                                int16_t num_rows, int16_t num_columns, int32_t num_rows_columns, int8_t formerNum, int16_t tailLength, 
                                TPipe* pipeIn) {
        this->num_rows = num_rows;
        this->num_columns = num_columns;
        this->num_rows_columns = num_rows_columns;

        if (GetBlockIdx() < formerNum)
        {
            this->batch_shape_size = tailLength + 1;
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y + this->num_rows_columns * this->batch_shape_size * GetBlockIdx());
        }
        else
        {
            this->batch_shape_size = tailLength;
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y +  this->num_rows_columns * this->batch_shape_size * GetBlockIdx() + this->num_rows_columns * formerNum);
        }

        pipeIn->InitBuffer(outQueueY, BUFFER_NUM, this->batch_shape_size * 32 + 256);
    }
    __aicore__ inline void Process() {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        if constexpr (std::is_same_v<DTYPE_Y, double>)
        {
            Duplicate(yLocal.ReinterpretCast<uint32_t>(), (uint32_t)0, this->batch_shape_size*8);

            uint64_t mask[2] = { 0xAAAAAAAAAAAAAAAA, 0xAAAAAAAAAAAAAAAA};
            Duplicate(yLocal.ReinterpretCast<uint32_t>(), (uint32_t)0x3FF00000, mask, (this->batch_shape_size/8 + 1), 1, 8 );
        }
        else if constexpr (std::is_same_v<DTYPE_Y, half>)
        {
            Duplicate(yLocal, (DTYPE_Y)1, this->batch_shape_size*16);
        }
        else
        {
            Duplicate(yLocal, (DTYPE_Y)1, this->batch_shape_size*8);
        }
        
        outQueueY.EnQue<DTYPE_Y>(yLocal);
        yLocal = outQueueY.DeQue<DTYPE_Y>();

        DataCopyExtParams copyParams{static_cast<uint16_t>(this->batch_shape_size), 1 * sizeof(DTYPE_Y), 0, static_cast<uint32_t>((this->num_rows_columns-1) * sizeof(DTYPE_Y)), 0};
        for(int32_t k=0; k<this->num_rows; k++)
        {
            DataCopyPad(yGm[k*(this->num_columns+1)], yLocal, copyParams);
        }
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe* pipe;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;

    GlobalTensor<DTYPE_Y> yGm;
    int16_t num_rows;
    int16_t num_columns;
    int32_t num_rows_columns;
    int32_t batch_shape_size;
    int16_t interval;
};

extern "C" __global__ __aicore__ void eye(GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
        TPipe pipe;
        if constexpr (std::is_same_v<DTYPE_Y, double>)
        {
            KernelEye op;
            op.Init(y,
                    tiling_data.num_rows, tiling_data.num_columns, tiling_data.num_rows_columns, tiling_data.formerNum, tiling_data.tailLength, 
                    &pipe);  
            op.Process();
        }
        else
        {
            KernelEye_all op;
            op.Init(y,
                    tiling_data.num_rows, tiling_data.num_columns, tiling_data.num_rows_columns, tiling_data.formerNum, tiling_data.tailLength, 
                    &pipe);  
            op.Process();
        }
}